/**
 * Copyright © 2013 Creativor Studio.
 * About license information, please see LICENSE.txt.
 */
package org.creativor.rayson.share.transport.nio;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

import org.creativor.rayson.share.annotation.ThreadSafe;
import org.creativor.rayson.util.LangUtil;

/**
 * An data buffer or NIO writing.
 * 
 * @author creativor
 */
@ThreadSafe
public class WriteBuffer {

	private final ByteBuffer bb;
	private final SocketChannel channel;
	private final ReentrantLock lock;
	private final Condition writeOut;
	private int channelWriteIndex;
	private final SelectionKey selectionKey;
	private final int capacity;

	public WriteBuffer(final int capacity, final SocketChannel channel,
			final SelectionKey selectionKey) {
		if (capacity <= 0)
			throw new IllegalArgumentException(
					"capacity should not less than 0");
		LangUtil.testNullArgument(channel, "Socket channel");
		LangUtil.testNullArgument(selectionKey, "Selection key");

		this.channel = channel;
		this.selectionKey = selectionKey;
		this.capacity = capacity;
		bb = ByteBuffer.allocate(capacity);
		lock = new ReentrantLock();
		writeOut = lock.newCondition();
		channelWriteIndex = 0;
	}

	/**
	 * Asynchrony write buffered data to underling socket channel.
	 * 
	 * @return The number of bytes written, possibly zero.
	 * @throws IOException
	 *             If failed to write to channel.
	 */
	public int asyncWrite() throws IOException {
		lock.lock();
		try {
			final int oldPos = bb.position();
			if (oldPos == 0)
				return 0;
			bb.flip();
			bb.position(channelWriteIndex);
			final int count = channel.write(bb);
			if (count > 0) {
				channelWriteIndex += count;
				if (channelWriteIndex == oldPos) {
					// un-register write event.
					this.selectionKey.interestOps(SelectionKey.OP_READ);
				}

				bb.position(oldPos);
				bb.limit(bb.capacity());
				writeOut.signalAll();
			}
			return count;
		} finally {
			lock.unlock();
		}
	}

	/**
	 * Flush out the buffered data.
	 * 
	 * @throws InterruptedException
	 *             If current thread is interrupted.
	 */
	public void flush() throws InterruptedException {

		lock.lock();
		try {
			final int pos = bb.position();
			if (pos == 0)
				return;
			// Register write key.
			this.selectionKey.interestOps(this.selectionKey.interestOps()
					| SelectionKey.OP_WRITE);
			this.selectionKey.selector().wakeup();

			while (channelWriteIndex < bb.position()) {
				writeOut.await();
			}
			// reset the buffer.
			channelWriteIndex = 0;
			bb.clear();
		} finally {
			lock.unlock();
		}

	}

	/**
	 * @return the capacity
	 */
	public int getCapacity() {
		return capacity;
	}

	/**
	 * Writes a sequence of bytes to this buffer.
	 * 
	 * @throws InterruptedException
	 * 
	 */
	public void write(final byte[] data) throws InterruptedException {
		if (data == null || data.length == 0)
			return;
		int index = 0;
		final int total = data.length;
		int bbRemaining;
		int leftCount;
		int writeCount;
		lock.lock();
		try {
			while ((leftCount = (total - index)) > 0) {
				bbRemaining = bb.remaining();
				if (bbRemaining == 0) {
					flush();
					continue;
				}
				writeCount = (leftCount > bbRemaining) ? bbRemaining
						: leftCount;
				while (writeCount > 0) {
					bb.put(data[index++]);
					writeCount--;
				}
			}
		} finally {
			lock.unlock();
		}
	}
}